import csv
import numpy as np
import sys

"""
Load data from file

:param 
-------
num_class: number of class

:method
-------
1. transform txt to csv
    using shell script:  cat *.txt | tr '\t' ','-> *.csv
2. read data from csv files

:return
-------
    labels and X
"""


class dataset(object):
    @staticmethod
    def load_data(num_class):
        assert num_class in [3, 4, 6, 8], 'num_class must be 3, 4, 6 or 8'
        labels, X = [], []
        file = './dataset/GMM%d.csv' % num_class
        with open(file) as f:
            f_csv = csv.DictReader(f)
            for row in f_csv:
                ret = list(map(eval, list(row.values())))
                labels.append(ret[0])
                X.append(ret[1:])
        return np.array(labels), np.array(X)


if __name__ == '__main__':
    num_class = 6
    if len(sys.argv) == 2:
        num_class = int(sys.argv[1])
    a, b = dataset.load_data(num_class)
    # a, b = load_data(num_class)
    print(a[0], '\n', b[0, :])
